import torch
from torchvision import transforms
from PIL import Image
import argparse
import json
import sys
import os

# 将项目根目录添加到模块搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.network import FontModel  # 导入模型定义

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载类别顺序
def load_class_names(config_path="class_names.json"):
    """从配置文件中加载类别顺序"""
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"配置文件 {config_path} 未找到，请确保训练脚本已生成该文件")
    with open(config_path, "r", encoding="utf-8") as f:
        class_names = json.load(f)
    return class_names

# 推理函数
def predict(image_path, model_path="models/font_model.pth", config_path="class_names.json"):
    # 加载类别顺序
    class_names = load_class_names(config_path)
    
    # 加载模型
    model = FontModel(num_classes=len(class_names))  # 类别数需与训练时一致
    model.load_state_dict(torch.load(model_path, weights_only=True))
    model.eval()
    
    # 加载并预处理图像
    image = Image.open(image_path).convert("RGB")
    image = transform(image).unsqueeze(0)  # 添加batch维度
    
    # 推理
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output, 1)
    return class_names[predicted.item()]

# 遍历文件夹并预测
def predict_folder(folder_path, model_path="models/font_model.pth", config_path="class_names.json"):
    # 支持的图片格式
    supported_formats = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".gif"]
    
    # 加载类别顺序
    class_names = load_class_names(config_path)

    # 遍历文件夹
    for root, _, files in os.walk(folder_path):
        for file in files:
            # 检查文件格式
            if any(file.lower().endswith(format) for format in supported_formats):
                image_path = os.path.join(root, file)
                try:
                    # 预测
                    result = predict(image_path, model_path, config_path)
                    print(f"图片: {image_path}, 预测字体: {result}")
                except Exception as e:
                    print(f"图片 {image_path} 预测失败: {e}")

if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=str, help="输入图片路径")
    parser.add_argument("--folder", type=str, help="输入文件夹路径")
    parser.add_argument("--model", type=str, default="models/font_model.pth", help="模型权重路径")
    parser.add_argument("--config", type=str, default="class_names.json", help="类别顺序配置文件路径")
    args = parser.parse_args()
    
    # 根据参数调用推理函数
    if args.image:
        result = predict(args.image, args.model, args.config)
        print(f"预测字体: {result}")
    elif args.folder:
        predict_folder(args.folder, args.model, args.config)
    else:
        print("请提供 --image 或 --folder 参数")